import numpy as np
from scipy.interpolate import RegularGridInterpolator

from GenericHdf import GenericHdf
from utils import find_files, random_value


class Fy3dL2VsmHdf(GenericHdf):

    def __init__(self, file_path):
        super().__init__(file_path)
        self.VSM_A = self.file["VSM_A"]
        self.VSM_D = self.file["VSM_D"]
        self.vsm_a_data = Fy3dL2VsmHdf.par(self.VSM_A)
        self.vsm_d_data = Fy3dL2VsmHdf.par(self.VSM_D)
        shape = self.vsm_a_data.shape
        self.geo_transform = [-180, 360.0 / shape[1], 0, 90, 0, -180.0 / shape[0]]
        self.projection = 'GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563,AUTHORITY["EPSG","7030"]],' \
                          'AUTHORITY["EPSG","6326"]],PRIMEM["Greenwich",0,AUTHORITY["EPSG","8901"]],' \
                          'UNIT["degree",0.0174532925199433,AUTHORITY["EPSG","9122"]],' \
                          'AUTHORITY["EPSG","4326"]]'
        pass

    @classmethod
    def par(cls, dataset):
        fill_value = dataset.attrs.get("FillValue", -999)
        intercept = dataset.attrs.get("Intercept", 0.0)
        slope = dataset.attrs.get("Slope", 0.001)
        data = np.array(dataset).astype(np.float32)
        indices = np.where(data == fill_value)
        data[indices] = np.nan
        return data * slope + intercept

    # 将目标经纬度转为行列坐标，返回浮点数(x, y)，由小数部分表示匹配偏差（后续插值用）
    def latlon_to_xy(self, lat, lon):
        global_rows, global_cols = self.VSM_A.shape
        x = (lon + 180.) / 360. * global_cols
        y = (lat + 90.) / 180. * global_rows
        return y, x

    # 求行列坐标的观测数据，行列坐标为浮点数(x, y)，根据坐标值的小数部分对每一层的数据进行插值运算
    def vsm_a_at_xy(self, x, y):
        global_rows, global_cols = self.VSM_A.shape
        interpolator = RegularGridInterpolator((np.arange(global_rows), np.arange(global_cols)), self.vsm_a_data)
        return interpolator((y, x))

    def vsm_d_at_xy(self, x, y):
        global_rows, global_cols = self.VSM_A.shape
        interpolator = RegularGridInterpolator((np.arange(global_rows), np.arange(global_cols)), self.vsm_d_data)
        return interpolator((y, x))

    def vsm_a_at(self, lat, lon):
        global_rows, global_cols = self.VSM_A.shape
        x = (lon + 180.) / 360. * global_cols
        y = (lat + 90.) / 180. * global_rows
        interpolator = RegularGridInterpolator((np.arange(global_rows), np.arange(global_cols)), self.vsm_a_data)
        return interpolator((y, x))

    def vsm_d_at(self, lat, lon):
        global_rows, global_cols = self.VSM_A.shape
        x = (lon + 180.) / 360. * global_cols
        y = (lat + 90.) / 180. * global_rows
        interpolator = RegularGridInterpolator((np.arange(global_rows), np.arange(global_cols)), self.vsm_d_data)
        return interpolator((y, x))

    def mock_vsm_a_at(self, lat, lon):
        value = self.vsm_a_at(lat, lon)
        if np.isnan(value):
            value = random_value(self.vsm_a_data)
        return value

    def mock_vsm_d_at(self, lat, lon):
        value = self.vsm_a_at(lat, lon)
        if np.isnan(value):
            value = random_value(self.vsm_d_data)
        return value


class Fy3dL2VsmHdfBatch:
    def __init__(self, fy3d_vsm_dir):
        file_paths = find_files(fy3d_vsm_dir, lambda f: '_L2_VSM_' in f.upper() and f.upper().endswith('.HDF'))
        self.fy3d_vsm_hdfs = []
        for file_path in file_paths:
            self.fy3d_vsm_hdfs.append(Fy3dL2VsmHdf(file_path))
        self.vsm_a_data = np.array([fy3d_vsm_hdf.vsm_a_data for fy3d_vsm_hdf in self.fy3d_vsm_hdfs])
        self.vsm_d_data = np.array([fy3d_vsm_hdf.vsm_d_data for fy3d_vsm_hdf in self.fy3d_vsm_hdfs])
        pass

    def max_vsm_a_at_mn(self, m, n):
        return np.nanmax(self.vsm_a_data[:, m, n])

    def mean_vsm_a_at_mn(self, m, n):
        return np.nanmax(self.vsm_a_data[:, m, n])

    def max_vsm_d_at_mn(self, m, n):
        return np.nanmean(self.vsm_d_data[:, m, n])

    def mean_vsm_d_at_mn(self, m, n):
        return np.nanmean(self.vsm_d_data[:, m, n])
